#ifndef SHERIFFOS2_SNAIL_MOTOR_H
#define SHERIFFOS2_SNAIL_MOTOR_H

#include <sheriffos.h>
#include <units.h>

constexpr float kSnailAnglePerRound = 360.0F;
constexpr float kEncoderFullScale = 65536.0F;
constexpr float kPWMMinDuty = 0.5F;
constexpr float kPWMMaxDuty = 0.9F;
constexpr float kPWMRangeScale = 0.4F / 50.0F;

struct SnailEncoderData {
    units::angle::degree_t rawAngle{};
    units::angle::degree_t lastAngle{};
    int32_t turnCount{};
    int32_t lastTurnCount{};
    float speed_radps{};
    float linearVelocity{};
    units::time::microsecond_t lastUpdateTime{};

    float diameterMeters{};
};

class SnailEncoder {
    SPI_HandleTypeDef* hspi_ = nullptr;
    GPIO_TypeDef* csPort_ = nullptr;
    uint16_t csPin_ = 0;

    uint8_t rxBuffer_[4]{};
    SnailEncoderData data_{};

public:
    SnailEncoder() = default;

    SnailEncoder(SPI_HandleTypeDef* hspi, GPIO_TypeDef* csPort, uint16_t csPin, float diameter)
        : hspi_(hspi), csPort_(csPort), csPin_(csPin) {
        data_.diameterMeters = diameter;
    }

    void setDiameter(float diameter) { data_.diameterMeters = diameter; }

    [[nodiscard]] SnailEncoderData const& getData() const { return data_; }

    void init() {
        uint8_t tx[4]{0x00};
        for (int i = 0; i < 8; i++) {
            tx[0] = 0x00;
            tx[1] = 0x00;
            HAL_GPIO_WritePin(csPort_, csPin_, GPIO_PIN_RESET);
            HAL_SPI_TransmitReceive(hspi_, tx, rxBuffer_, 2, 100);
            HAL_GPIO_WritePin(csPort_, csPin_, GPIO_PIN_SET);
            osDelay(1);
        }
    }

    void readRaw() {
        uint8_t tx[4]{0};
        HAL_GPIO_WritePin(csPort_, csPin_, GPIO_PIN_RESET);
        HAL_SPI_TransmitReceive(hspi_, tx, rxBuffer_, 4, 100);
        HAL_GPIO_WritePin(csPort_, csPin_, GPIO_PIN_SET);
        this->decode(rxBuffer_);
    }

    void decode(uint8_t const* buffer) {
        auto now_time_us = bsp::getGlobalTime_us_i64();

        auto lastAngle = data_.rawAngle;
        data_.lastAngle = lastAngle;
        data_.lastTurnCount = data_.turnCount;

        auto rawEncoder = static_cast<int16_t>((buffer[0] << 8) | buffer[1]);
        data_.rawAngle = units::make_unit<units::angle::degree_t>(
            static_cast<float>(rawEncoder) * kSnailAnglePerRound / kEncoderFullScale - 180.0F);

        auto diff = (data_.rawAngle - data_.lastAngle).to<float>();

        // 防过零跳变修正
        if (diff > 180.0F)
            diff -= 360.0F;
        else if (diff < -180.0F)
            diff += 360.0F;

        if (diff > 0 && data_.lastAngle.to<float>() < -90 && data_.rawAngle.to<float>() > 90)
            ++data_.turnCount;
        else if (diff < 0 && data_.lastAngle.to<float>() > 90 && data_.rawAngle.to<float>() < -90)
            --data_.turnCount;

        auto dt = static_cast<float>(now_time_us - data_.lastUpdateTime.to<int64_t>()) / 1'000'000.0f;
        data_.lastUpdateTime = units::make_unit<units::time::microsecond_t>(now_time_us);

        data_.speed_radps = (diff / 180.0F * 3.1415926F) / dt;
        data_.linearVelocity = data_.speed_radps * data_.diameterMeters * 0.5F;
    }
};

class SnailMotor final {
    TIM_HandleTypeDef* htim_ = nullptr;
    uint32_t pwmChannel_ = 0;

    float output_ = 0.0F;
    SnailEncoder& encoder_;

public:
    SnailMotor() = delete;

    SnailMotor(TIM_HandleTypeDef* htim, uint32_t channel, SnailEncoder& encoder)
        : htim_(htim), pwmChannel_(channel), encoder_(encoder) {}

    void init(bool encoderInit = true) {
        if (encoderInit) encoder_.init();
        HAL_TIM_PWM_Start(htim_, pwmChannel_);
    }

    void setOutput(float val) { output_ = val; }

    float getOutput() const { return output_; }

    void sendOutput() {
        float duty = output_ * kPWMRangeScale + 0.5F;
        duty = std::clamp(duty, kPWMMinDuty, kPWMMaxDuty);
        auto pulse = duty * static_cast<float>(htim_->Init.Period + 1);
        __HAL_TIM_SetCompare(htim_, pwmChannel_, static_cast<uint32_t>(pulse));
    }

    void readEncoder() { encoder_.readRaw(); }

    [[nodiscard]] SnailEncoder& getEncoder() { return encoder_; }

    [[nodiscard]] SnailEncoderData const& getEncoderData() const { return encoder_.getData(); }
};
#endif  // SHERIFFOS2_SNAIL_MOTOR_H